#include "SymmetricEncrypter.h"
#include <botan/botan.h>
#include <botan/algo_factory.h>
#include <memory>
#include <assert.h>
#include <fstream>
#include <algorithm>

namespace kryptos { namespace cryptography {

struct SymmetricEncrypter::Impl
{
    SymmetricEncrypter&         outer;
    Botan::Keyed_Filter*        currentAlgo;
    Botan::SymmetricKey         symkey;
    Botan::InitializationVector iv;

    enum OpType { ENCRYPT, DECRYPT};

    Impl(SymmetricEncrypter& encrypter)
        : outer(encrypter)
    {
    }

    std::auto_ptr<Botan::Pipe> getPipe(OpType type, bool encode)
    {
        std::auto_ptr<Botan::Pipe> pipe;

        if (type == ENCRYPT)
        {
            pipe.reset(new Botan::Pipe(currentAlgo = Botan::get_cipher("AES-256/CBC/PKCS7", symkey, iv, Botan::ENCRYPTION)));
            if (encode)
                pipe->append(new Botan::Base64_Encoder(true, 80, true));
        }
        else
        {
            pipe.reset(new Botan::Pipe(currentAlgo = Botan::get_cipher("AES-256/CBC/PKCS7", symkey, iv, Botan::DECRYPTION)));
            if (encode)
                pipe->append(new Botan::Base64_Decoder());
        }

        return pipe;
    }

    std::auto_ptr<Botan::Pipe> getPipeWithSink(OpType type, std::ostream &synkFile, bool encode)
    {
        std::auto_ptr<Botan::Pipe> pipe = getPipe(type, encode);
        pipe->append(new Botan::DataSink_Stream(synkFile));

        return pipe;
    }


    // Buffer encryption/decryption
    int processBuffer(OpType             type,
                      const char*        buffer,
                      int                len,
                      std::vector<char>& output,
                      bool               encode=false)
    {
        std::auto_ptr<Botan::Pipe> pipe = getPipe(type, encode);

        symkey.change((Botan::byte*)outer.m_key.nextBytes(outer.m_keyBlockSize).data(), outer.m_keyBlockSize);
        currentAlgo->set_key(symkey);

        pipe->process_msg((Botan::byte*)buffer, len);
        Botan::SecureVector<Botan::byte> vec = pipe->read_all(0);
        output.clear();
        output.resize(vec.size());
        std::copy(vec.begin(), vec.end(), output.begin());

        return 0;
    }

    // File encryption/decryption
    int processFile(OpType             type,
                    std::istream&      input,
                    unsigned long long cipherEnd,
                    std::ostream&      output,
                    bool               encode=false)
    {
        if (input.good() && !input.eof())
        {
            std::auto_ptr<Botan::Pipe> pipe = getPipeWithSink(type, output, encode);
            doProcess(type, input, cipherEnd, *pipe.get());
        }

        return 0;
    }

    int doProcess(OpType             type,
                  std::istream&      input,
                  unsigned long long cipherEnd,
                  Botan::Pipe&       pipe)
    {
        pipe.start_msg();

        unsigned long long currentPos = input.tellg();
        unsigned long long cryptoBytesLeft = cipherEnd - currentPos;

        while (!input.eof() && cryptoBytesLeft > 0)
        {
            // TODO: bufferizar la lectura, leer de a 4096 bytes, y luego en memoria
            // procesar los bloques de 32bytes, cambiando la clave simetrica
            symkey.change((Botan::byte*)outer.m_key.nextBytes(outer.m_keyBlockSize).data(), outer.m_keyBlockSize);
            currentAlgo->set_key(symkey);

            int sizeToRead = outer.m_keyBlockSize > cryptoBytesLeft ? cryptoBytesLeft : outer.m_keyBlockSize;

            std::vector<char> buffer(sizeToRead);
            input.read(buffer.data(), sizeToRead);
            buffer.resize(input.gcount());
            pipe.write((Botan::byte*)buffer.data(), buffer.size());

            cryptoBytesLeft -= buffer.size();
        }

        pipe.end_msg();

        return 0;
    }
};

/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////

SymmetricEncrypter::SymmetricEncrypter(SymmetricKey& key)
    : m_key(key)
    , m_impl(new Impl(*this))
{
    std::auto_ptr<Botan::BlockCipher> blockCipher(Botan::get_block_cipher("AES-256"));

    m_keyBlockSize = blockCipher->key_spec().minimum_keylength();
    m_blockSize    = blockCipher->block_size();

    m_impl->symkey.change((Botan::byte*)m_key.nextBytes(m_keyBlockSize).data(), m_keyBlockSize);

    // moves the offste back since we are calling nextBytes again on each encrypt/decrypt action
    m_key.goBack(m_keyBlockSize);

    std::vector<char> zeroedIV(m_blockSize);
    memset((char*)&zeroedIV[0], 0, m_blockSize);
    m_impl->iv.change((const unsigned char*)zeroedIV.data(), m_blockSize);
}

SymmetricEncrypter::~SymmetricEncrypter()
{
}

int SymmetricEncrypter::encryptFile(const std::string& filepath, const std::string& output)
{
    assert(m_key.isValid());

    std::ifstream inputFile(filepath.c_str(), std::ios::binary);
    std::ofstream outputFile(output.c_str(), std::ios::binary);

    inputFile.seekg(0, std::ios::end);
    unsigned long long inputEnd = inputFile.tellg();
    inputFile.seekg(0);

    // TODO: Check for input/output file and filesize with keysize
    m_impl->processFile(Impl::ENCRYPT, inputFile, inputEnd, outputFile);

    return 0;
}

int SymmetricEncrypter::encryptFile(const std::string &filepath, std::ostream &output)
{
    assert(m_key.isValid());

    std::ifstream inputFile(filepath.c_str(), std::ios::binary);
    inputFile.seekg(0, std::ios::end);
    unsigned long long inputEnd = inputFile.tellg();
    inputFile.seekg(0);

    m_impl->processFile(Impl::ENCRYPT, inputFile, inputEnd, output);

    return 0;
}

int SymmetricEncrypter::encryptAndEncodeFile(const std::string &filepath, std::ostream &output)
{
    assert(m_key.isValid());

    std::ifstream inputFile(filepath.c_str(), std::ios::binary);

    inputFile.seekg(0, std::ios::end);
    unsigned long long inputEnd = inputFile.tellg();
    inputFile.seekg(0);

    m_impl->processFile(Impl::ENCRYPT, inputFile, inputEnd, output, true);

    return 0;
}

int SymmetricEncrypter::encryptBuffer(const char* buffer, int len, std::vector<char>& output)
{
    assert(buffer);
    assert(len > 0);
    assert(m_key.isValid());

    m_impl->processBuffer(Impl::ENCRYPT, buffer, len, output);

    return 0;
}

int SymmetricEncrypter::encryptAndEncodeBuffer(const char* buffer, int len, std::vector<char> &output)
{
    assert(buffer);
    assert(len > 0);
    assert(m_key.isValid());

    m_impl->processBuffer(Impl::ENCRYPT, buffer, len, output, true);

    return 0;
}

int SymmetricEncrypter::decryptFile(const std::string& filepath, const std::string& output)
{
    assert(m_key.isValid());

    std::ifstream inputFile(filepath.c_str(), std::ios::binary);
    std::ofstream outputFile(output.c_str(), std::ios::binary);

    inputFile.seekg(0, std::ios::end);
    unsigned long long inputEnd = inputFile.tellg();
    inputFile.seekg(0);

    //TODO: Check for input/output file and filesize with keysize
    m_impl->processFile(Impl::DECRYPT, inputFile, inputEnd, outputFile);

    return 0;
}

int SymmetricEncrypter::decryptFile(std::istream &inputfile, std::ostream &output)
{
    assert(m_key.isValid());

    unsigned long long current = inputfile.tellg();
    inputfile.seekg(0, std::ios::end);
    unsigned long long inputEnd = inputfile.tellg();
    inputfile.seekg(current);

    m_impl->processFile(Impl::DECRYPT, inputfile, inputEnd, output);

    return 0;
}

int SymmetricEncrypter::decryptFile(std::istream &inputfile, unsigned long long cipherEnd, std::ostream &output)
{
    assert(m_key.isValid());
    m_impl->processFile(Impl::DECRYPT, inputfile, cipherEnd, output);

    return 0;
}

int SymmetricEncrypter::decryptBuffer(const char* buffer, int len, std::vector<char>& output)
{
    assert(buffer);
    assert(len > 0);
    assert(m_key.isValid());

    m_impl->processBuffer(Impl::DECRYPT, buffer, len, output);

    return 0;
}

}}
